# -*- coding: utf-8 -*-

import torch.nn as nn
import torch.nn.functional as F


class RegionProposalNetwork(nn.Module):
    def __init__(self,
                 in_channel=2048,
                 mid_channel=512,
                 n_base_anchors=9):
        super(RegionProposalNetwork, self).__init__()

        self.n_base_anchors = n_base_anchors
        self.conv1 = nn.Conv2d(in_channel, mid_channel, 3, 1, 1, bias=False)
        self.bn1 = nn.BatchNorm2d(mid_channel)
        self.relu1 = nn.ReLU(inplace=True)
        self.object_score = nn.Conv2d(mid_channel, self.n_base_anchors * 2, 1, 1, 0)
        self.object_location = nn.Conv2d(mid_channel, self.n_base_anchors * 4, 1, 1, 0)

        self._weight_init()

    def _weight_init(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        object_score = self.object_score(x)
        object_location = self.object_location(x)
        object_score = object_score.permute(0, 2, 3, 1).contiguous().view(-1, 2)
        object_location = object_location.permute(0, 2, 3, 1).contiguous().view(-1, 4)
        return object_score, object_location


def compute_rpn_loss(predict_scores, predict_locations, object_labels, location_targets):
    # location loss
    pos_predict_locations = predict_locations[object_labels == 1]
    pos_location_targets = location_targets[object_labels == 1]
    location_loss = F.smooth_l1_loss(
        pos_predict_locations, pos_location_targets, reduction='sum')
    # classification loss
    valid_predict_scores = predict_scores[object_labels >= 0]
    valid_object_labels = object_labels[object_labels >= 0]
    classification_loss = F.nll_loss(
        F.log_softmax(valid_predict_scores, dim=1),
        valid_object_labels, reduction='sum')
    rpn_loss = location_loss + classification_loss
    return rpn_loss, location_loss, classification_loss
